import json  # 用于记录完整session
from openai import OpenAI
import pymysql

def get_completion(prompt):

    conn = pymysql.connect(
        host='localhost',
        port=3306,
        user='root',
        passwd='123456',
        database='datas'
    )

    SQL = []

    cursor = conn.cursor()
    cursor.execute('SHOW TABLES')
    tables = cursor.fetchall()
    for table in tables:
        SQL.append(table[0])
    i = 0
    while i < len(SQL):
        cursor.execute(f'SHOW COLUMNS FROM {SQL[i]}')
        columns = cursor.fetchall()
        SQL[i] += '('
        for column in columns:
            SQL[i] += f'{column[0]},'
        SQL[i] = SQL[i][:-1]
        SQL[i] += ')'
        i += 1


    client = OpenAI(
        api_key="sk-RvVVNjgBjYGxsD4hB4LXwIxWI4sEyz8hhldUEaZQ9fFjUQSV",
        base_url="https://api.fe8.cn/v1",
    )

    SQLData = ''
    for item in SQL:
        SQLData += f'{item} '

    session = [
        {
            "role": "system",
            "content": "你需要结合给出的自然语言给出用户所需要的SQL语句，不要说任何一个多余的字，只需要给出SQL语句，数据库内的表与其中列名如下，其格式为表1(列1,列2,列3...)：{} ".format(SQLData),
        }
    ]


#def get_completion(prompt):
    session.append({"role": "user", "content": prompt})  # 在会话中加入user的prompt

    response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=session
    )  # 基于会话，获取assistant的回复
    # 注意：此部分可传入诸多参数，通过它们来控制assistant回复的规制，详情可见。

    msg = response.choices[0].message.content  # 获取assistant的回复
    session.append({"role": "assistant", "content": msg})  # 在会话中加入assistant的回复
    #return msg
    cursor = conn.cursor()
    cursor.execute(msg)
    print(msg)
    return cursor.fetchall()

#if __name__ == "__main__":
#    #get_completion(prompt="查询学生年龄")
#     #print(json.dumps(session, indent=4, ensure_ascii=False))
    #print(get_completion(prompt="查询年级一样的老师和学生"))